#!/usr/bin/env perl

# Copyright (C) Yichun Zhang (agentzh)
# Copyright (C) Guanlan Dai

# TODO: port this script into the nginx core for greater flexibility
# and better performance.

use strict;
use warnings;

our $VERSION = 0.01;

use File::Spec ();
use FindBin ();
use List::Util qw( max );
use Getopt::Long qw( GetOptions :config no_ignore_case);
use File::Temp qw( tempdir );
use POSIX qw( WNOHANG );

my @all_args =  @ARGV;
GetOptions("c=i",           \(my $conns_num),
           "e=s",           \(my $src),
           "h|help",        \(my $help),
           "I=s@",          \(my $Inc),
           "nginx=s",       \(my $nginx_path),
           "valgrind",      \(my $use_valgrind),

           "valgrind-opts=s", \(my $valgrind_opts),

           "V",             \(my $version))
   or die usage(1);

if ($help) {
    usage(0);
}

if (!$nginx_path) {
    $nginx_path = File::Spec->catfile($FindBin::Bin, "..", "nginx", "sbin", "nginx");
    if (!-f $nginx_path) {
        $nginx_path = "nginx";  # find in PATH
    }
}

#warn $nginx_path;

if ($version) {
    warn "resty $VERSION\n";
    my $cmd = "$nginx_path -V";
    exec $cmd or die "Failed to run command \"$cmd\": $!\n";
}

my $lua_package_path_config = '';
if (defined $Inc) {
    my $package_path = "";
    my $package_cpath = "";
    for my $dir (@$Inc) {
        if (!-d $dir) {
            die "Search directory $dir is not found.\n";
        }
        $package_path .= File::Spec->catfile($dir, "?.lua;");
        $package_cpath .= File::Spec->catfile($dir, "?.so;");
    }
    $lua_package_path_config = <<_EOC_;
    lua_package_path "$package_path;";
    lua_package_cpath "$package_cpath;";
_EOC_
}

my $luafile;
if (!defined $src) {
    $luafile = shift
        or die qq{No Lua input file nor -e "" option specified.\n};
}

my $conns = $conns_num || 64;

my @nameservers;

# try to read the nameservers used by the system resolver:
if (open my $in, "/etc/resolv.conf") {
    while (<$in>) {
        if (/^\s*nameserver\s+(\d+(?:\.\d+){3})(?:\s+|$)/) {
            push @nameservers, $1;
            if (@nameservers > 10) {
                last;
            }
        }
    }
    close $in;
}

if (!@nameservers) {
    # default to Google's open DNS servers
    push @nameservers, "8.8.8.8", "8.8.4.4";
}

#warn "@nameservers\n";

my $prefix_dir = tempdir(CLEANUP => 1);
#warn "prefix dir: $prefix_dir\n";

my $logs_dir = File::Spec->catfile($prefix_dir, "logs");
mkdir $logs_dir or die "failed to mkdir $logs_dir: $!";

my $conf_dir = File::Spec->catfile($prefix_dir, "conf");
mkdir $conf_dir or die "failed to mkdir $conf_dir: $!";

my $chunk_name;
if (defined $src) {
    $luafile = File::Spec->catfile($conf_dir, "a.lua");
    open my $out, ">$luafile"
        or die "Cannot open $luafile for writing: $!\n";
    print $out $src;
    close $out;
    $chunk_name = "=(command line -e)";

} else {
    $chunk_name = "\@$luafile";
}

my $quoted_luafile = quote_as_lua_str($luafile);

my @user_args =  @ARGV;
my $args = gen_lua_code_for_args(\@user_args, \@all_args);

my $loader = <<_EOC_;
            local gen
            do
                $args
                local fname = $quoted_luafile
                local f = assert(io.open(fname, "r"))
                local chunk = f:read("*a")
                gen = assert(loadstring(chunk, "$chunk_name"))
            end
_EOC_

my $env_list = '';
for my $var (sort keys %ENV) {
    #warn $var;
    $env_list .= "env $var;\n";
}

my $conf_file = File::Spec->catfile($conf_dir, "nginx.conf");
open my $out, ">$conf_file"
    or die "Cannot open $conf_file for writing: $!\n";

print $out <<_EOC_;
daemon off;
master_process off;
worker_processes 1;

$env_list

error_log /dev/stderr warn;
#error_log /dev/stderr debug;

events {
    worker_connections $conns;
}

http {
    access_log off;
    lua_socket_log_errors off;
    resolver @nameservers;
$lua_package_path_config
    init_by_lua '
        local stdout = io.stdout
        local ngx_null = ngx.null
        local maxn = table.maxn
        local unpack = unpack
        local concat = table.concat

        local expand_table
        function expand_table(src, inplace)
            local n = maxn(src)
            local dst = inplace and src or {}
            for i = 1, n do
                local arg = src[i]
                local typ = type(arg)
                if arg == nil then
                    dst[i] = "nil"

                elseif typ == "boolean" then
                    if arg then
                        dst[i] = "true"
                    else
                        dst[i] = "false"
                    end

                elseif arg == ngx_null then
                    dst[i] = "null"

                elseif typ == "table" then
                    dst[i] = expand_table(arg, false)

                elseif typ ~= "string" then
                    dst[i] = tostring(arg)

                else
                    dst[i] = arg
                end
            end
            return concat(dst)
        end

        local function output(...)
            local args = {...}

            return stdout:write(expand_table(args, true))
        end

        ngx.print = output
        ngx.say = function (...)
                local ok, err = output(...)
                if ok then
                    return output("\\\\n")
                end
                return ok, err
            end
        print = ngx.say

        ngx.flush = function (...) return stdout:flush() end
        -- we cannot close stdout here due to a bug in Lua:
        ngx.eof = function (...) return true end
        ngx.exit = os.exit
    ';

    init_worker_by_lua '
        local exit = os.exit
        local stderr = io.stderr

        local function handle_err(err)
            if err then
                err = string.gsub(err, "^init_worker_by_lua:%d+: ", "")
                stderr:write(err, "\\\\n")
            end
            return exit(1)
        end

        local ok, err = pcall(function ()
            if not ngx.config
               or not ngx.config.ngx_lua_version
               or ngx.config.ngx_lua_version < 9011
            then
                error("at least ngx_lua 0.9.12 is required")
            end

$loader
            -- print("calling timer.at...")
            local ok, err = ngx.timer.at(0, function ()
                -- io.stderr:write("timer firing")
                local ok, err = pcall(gen)
                if not ok then
                    return handle_err(err)
                end
                local rc = err
                if rc and type(rc) ~= "number" then
                    return handle_err("bad return value of type " .. type(rc))
                end
                return exit(rc)
            end)
            if not ok then
                return handle_err(err)
            end
            -- print("timer created")
        end)

        if not ok then
            return handle_err(err)
        end
    ';
}
_EOC_

close $out;

my $cmd = "$nginx_path -p $prefix_dir/ -c conf/nginx.conf";

if ($use_valgrind) {
    my $opts = '';
    if ($valgrind_opts) {
        $opts .= " $valgrind_opts";
    }
    $cmd = "valgrind$opts $cmd";
}

my $child_pid;

sub sigint {
    $SIG{INT} = \&sigint;
    if ($child_pid) {
        kill INT => $child_pid;
    }
}
$SIG{INT} = \&sigint;

my $pid = fork();

if (!defined $pid) {
    die "fork() failed: $!\n";
}

if ($pid == 0) {  # child process
    #warn "exec $cmd...";
    exec $cmd or die "Failed to run command \"$cmd\": $!\n";

} else {
    $child_pid = $pid;
    waitpid($child_pid, 0);
    my $rc = 0;
    if (defined $?) {
        $rc = ($? >> 8);
    }
    exit($rc);
}

sub usage {
    my $rc = shift;
    my $msg = <<_EOC_;
resty [options] [lua-file [args]]

Options:
    -c num      set maximal connection count (default: 64).
    -e prog     run the inlined Lua code in "prog".
    --help      print this help.
    -I dir      Add dir to the search paths for Lua libraries.
    --nginx     specify the nginx path (this option might be removed in the future).
    -V          print version numbers and nginx configurations.
    --valgrind  use valgrind to run the underyling nginx

    --valgrind-opts     pass extra options to valgrind

For bug reporting instructions, please see:
<http://openresty.org/#Community>
_EOC_
    if ($rc == 0) {
        print $msg;
        exit(0);
    }

    warn $msg;
    exit($rc);
}

sub get_bracket_level {
    my %bracket_levels;
    my $bracket_level = 0;
    my $max_level = 0;

    # scan all args and store level of closing brackets
    for my $arg (@_) {
        while ($arg =~ /\](=*)\]/g) {
            my $level = length($1);
            if ($level > $max_level) {
                $max_level = $level;
            }
            $bracket_levels{$level} = 1;
        }
    }

    # if args contain closing bracket
    if (%bracket_levels) {
        # find the shortest form of the long brackets accordingly
        for (my $i = 1; $i < $max_level; $i++) {
            if (!exists $bracket_levels{$i}) {
                $bracket_level = $i;
                last;
            }
        }

        if ($bracket_level == 0) {
            $bracket_level = $max_level + 1;
        }
        return $bracket_level;
    }
    return 1;
}

sub quote_as_lua_str {
    my ($str) = @_;
    my $bracket_level = get_bracket_level($str);
    my $left_bracket = "[" . "=" x $bracket_level . "[";
    my $right_bracket = "]" . "=" x $bracket_level . "]";

    return $left_bracket . $str . $right_bracket;
}

sub gen_lua_code_for_args {
    my ($user_args, $all_args) = @_;

    my $luasrc = "arg = {}\n";

    # args[n] (n = 0)
    $luasrc .= "arg[0] = $quoted_luafile\n";

    # args[n] (n > 0)
    for my $i (0 .. $#user_args) {
        my $index = $i + 1;
        my $quoted_arg = quote_as_lua_str($user_args[$i]);
        $luasrc .= "arg[$index] = $quoted_arg\n";
    }

    # args[n] (n < 0)
    my $left_num = $#all_args - $#user_args;
    for my $i (0 .. $left_num - 2) {
        my $index = 0 - $left_num + $i + 1;
        my $quoted_arg = quote_as_lua_str($all_args[$i]);
        $luasrc .= "arg[$index] = $quoted_arg\n";
    }

    # args[n] (n = the index of resty-cli itself)
    my $index = 0 - $left_num;
    my $quoted_arg = quote_as_lua_str($0);
    $luasrc .= "arg[$index] = $quoted_arg\n";

    #warn $luasrc;
    return $luasrc;
}
